#include "data_dumper.h"

#include <sys/stat.h>
#include <sys/types.h>
#include <unistd.h>
#include <numeric>
#include <fstream>
#include <map>
#include <vector>
#include <sstream>
#include <iterator>
#include <iostream>
#include "file_utils.h"

#define UNLIKELY(x) __builtin_expect(!!(x), 0)
#define LIKELY(x) __builtin_expect(!!(x), 1)

DataDumper::DataDumper() : path_(""), start_(false), init_(false) {}

DataDumper::~DataDumper() { UnInit(); }

DataDumper &DataDumper::GetInstance() {
  static DataDumper instance;
  return instance;
}

void DataDumper::Init(const std::string &path, int32_t rank_id,
                      size_t capacity) {
  path_ = path;
  rank_id_ = rank_id;
  data_chunk_buf_.Init(capacity);
  init_.store(true);
}

void DataDumper::UnInit() {
  if (init_.load()) {
    init_.store(false);
    start_.store(false);
    for (auto &f : fd_map_) {
      if (f.second != nullptr) {
        fclose(f.second);
        f.second = nullptr;
      }
    }
    fd_map_.clear();
  }
}

void DataDumper::Start() {
  if (!init_.load() || !Utils::CreateDir(path_)) {
    return;
  }
  start_.store(true);
}

void DataDumper::Stop() {
  if (start_.load() == true) {
    start_.store(false);
  }
  Flush();
}

void DataDumper::GatherAndDumpData() {
  std::map<std::string, std::vector<uint8_t>> dataMap;
  uint64_t batchSize = 0;
  while (batchSize < kBatchMaxLen) {
    std::unique_ptr<BaseReportData> data = data_chunk_buf_.Pop();
    if (data == nullptr) {
      break;
    }
    std::vector<uint8_t> encodeData = data->encode();
    batchSize += encodeData.size();
    const std::string &key = data->tag;
    auto iter = dataMap.find(key);
    if (iter == dataMap.end()) {
      dataMap.insert({key, encodeData});
    } else {
      iter->second.insert(iter->second.end(), encodeData.cbegin(),
                          encodeData.cend());
    }
  }
  if (dataMap.size() > 0) {
    Dump(dataMap);
  }
}

void DataDumper::Flush() {
  if (data_chunk_buf_.Size() == 0) {
    std::cout << "data_chunk_buf_ = " << data_chunk_buf_.Size() << std::endl;
    return;
  }
  uint64_t start_time = Utils::GetClockMonotonicRawNs();

  while (data_chunk_buf_.Size() > 0) {
    GatherAndDumpData();
  }
  data_chunk_buf_.Reset();
  uint64_t end_time = Utils::GetClockMonotonicRawNs();
  auto tid = Utils::GetTid();
  std::unique_ptr<OpRangeData> save_data = std::make_unique<OpRangeData>(
      start_time, end_time, tid, "Profiler_SaveDate", rank_id_);
  data_chunk_buf_.Push(std::move(save_data));
  GatherAndDumpData();
  data_chunk_buf_.Reset();
}

void DataDumper::Report(std::unique_ptr<BaseReportData> data) {
  if (!start_.load() || data == nullptr) {
    return;
  }
  int i = 0;
  while (is_flush_.load() && i < 10) {
    usleep(kMaxWaitTimeUs);
    i++;
  }
  if (!data_chunk_buf_.Push(std::move(data))) {
    is_flush_.store(true);
    std::lock_guard<std::mutex> flush_lock_(flush_mutex_);
    if (data_chunk_buf_.Full()) {
      Flush();
    }
    is_flush_.store(false);
    if (!data_chunk_buf_.Push(std::move(data))) {
    }
  }
}

void DataDumper::Dump(
    const std::map<std::string, std::vector<uint8_t>> &dataMap) {
  for (auto &data : dataMap) {
    FILE *fd = nullptr;
    const std::string dump_file = path_ + "/" + data.first;
    auto iter = fd_map_.find(dump_file);
    if (iter == fd_map_.end()) {
      if (!Utils::IsFileExist(dump_file) && !Utils::CreateDumpFile(dump_file)) {
        std::cout << "CreateDumpFile failed" << std::endl;
        continue;
      }
      fd = fopen(dump_file.c_str(), "ab");
      if (fd == nullptr) {
        std::cout << "fopen failed" << std::endl;
        continue;
      }
      fd_map_.insert({dump_file, fd});
    } else {
      fd = iter->second;
    }
    fwrite(reinterpret_cast<const char *>(data.second.data()), sizeof(char),
           data.second.size(), fd);
    fflush(fd);
  }
}


inline void EncodeStrData(uint16_t type, const std::string &data, const std::unique_ptr<std::vector<uint8_t>> &result) {
  for (size_t i = 0; i < sizeof(uint16_t); ++i) {
    result->push_back((type >> (i * 8)) & 0xff);
  }
  uint32_t length = data.size();
  for (size_t i = 0; i < sizeof(uint32_t); ++i) {
    result->push_back((length >> (i * 8)) & 0xff);
  }
  result->insert(result->end(), data.begin(), data.end());
}



void OpRangeData::preprocess() {
  const std::string delim = "|";
  const std::string remove_ms = "site-packages/mindspore";
  if (stack.size() > 0 && !stack[0].empty()) {
    std::string all_stack = stack[0];
    stack.erase(stack.begin());
    size_t nPos = all_stack.find(delim.c_str());
    while (nPos != std::string::npos) {
      std::string temp = all_stack.substr(0, nPos + 1);
      if (temp.find(remove_ms.c_str()) == std::string::npos) {
        stack.push_back(temp);
      }
      all_stack = all_stack.substr(nPos + 1);
      nPos = all_stack.find(delim.c_str());
    }
  }
}


template <typename T>
void EncodeFixedData(const std::vector<T> &data_list, const std::unique_ptr<std::vector<uint8_t>> &result) {
  for (auto data : data_list) {
    for (size_t i = 0; i < sizeof(T); ++i) {
      result->push_back((static_cast<size_t>(data) >> (i * 8)) & 0xff);
    }
  }
}

template <typename T>
void Encode2DIntegerMatrixDatas(const uint16_t type, const std::vector<std::vector<T>> &data_list,
                                const std::unique_ptr<std::vector<uint8_t>> &result) {
  std::string rst;
  for (auto tensor : data_list) {
    std::stringstream ss;
    copy(tensor.begin(), tensor.end(), std::ostream_iterator<T>(ss, ","));
    std::string str = ss.str();
    if (!str.empty()) {
      str.pop_back();
    }
    rst += (str + ";");
  }
  if (!rst.empty()) {
    rst.pop_back();
  }
  EncodeStrData(type, rst, result);
}

inline void EncodeStrArrayData(const uint16_t type, const std::vector<std::string> &data_list,
                               const std::unique_ptr<std::vector<uint8_t>> &result) {
  std::string rst = std::accumulate(data_list.begin(), data_list.end(), std::string(""),
                                    [](const std::string &r, const auto d) { return r + ';' + d; });
  if (!rst.empty()) {
    rst.pop_back();
  }
  EncodeStrData(type, rst, result);
}

inline void EncodeStrMapData(const uint16_t type, const std::map<std::string, std::string> &data_map,
                             const std::unique_ptr<std::vector<uint8_t>> &result) {
  std::string rst = std::accumulate(data_map.begin(), data_map.end(), std::string(""),
                                    [](const std::string &r, const std::pair<const std::string, std::string> &element) {
                                      return r + element.first + ":" + element.second + ";";
                                    });
  if (!rst.empty()) {
    rst.pop_back();
  }
  EncodeStrData(type, rst, result);
}

std::vector<uint8_t> OpRangeData::encode() {
  preprocess();
  std::unique_ptr<std::vector<uint8_t>> result = std::make_unique<std::vector<uint8_t>>();
  EncodeFixedData<int64_t>({start_ns, end_ns, sequence_number}, result);
  EncodeFixedData<uint64_t>({process_id, start_thread_id, end_thread_id, forward_thread_id, flow_id, step}, result);
  EncodeFixedData<int8_t>({level}, result);
  result->push_back(is_async);
  EncodeStrData(static_cast<uint16_t>(OpRangeDataType::NAME), name, result);
  if (!input_dtypes.empty()) {
    EncodeStrArrayData(static_cast<uint16_t>(OpRangeDataType::INPUT_DTYPES), input_dtypes, result);
  }
  if (!input_shapes.empty()) {
    Encode2DIntegerMatrixDatas<int64_t>(static_cast<uint16_t>(OpRangeDataType::INPUT_SHAPE), input_shapes, result);
  }
  if (!stack.empty()) {
    EncodeStrArrayData(static_cast<uint16_t>(OpRangeDataType::STACK), stack, result);
  }
  if (!custom_info.empty()) {
    EncodeStrMapData(static_cast<uint16_t>(OpRangeDataType::CUSTOM_INFO), custom_info, result);
  }
  if (!module_hierarchy.empty()) {
    EncodeStrArrayData(static_cast<uint16_t>(OpRangeDataType::MODULE_HIERARCHY), module_hierarchy, result);
  }
  std::vector<uint8_t> resultTLV;
  uint16_t dataType = static_cast<uint16_t>(OpRangeDataType::OP_RANGE_DATA);
  for (size_t i = 0; i < sizeof(uint16_t); ++i) {
    resultTLV.push_back((dataType >> (i * 8)) & 0xff);
  }
  uint32_t length = result->size();
  for (size_t i = 0; i < sizeof(uint32_t); ++i) {
    resultTLV.push_back((length >> (i * 8)) & 0xff);
  }
  resultTLV.insert(resultTLV.end(), result->cbegin(), result->cend());
  return resultTLV;
}